'''
Created on Sep 15, 2010

@author: oabalbin

This script parses the Small_RNA compare file
Look for miRNAs with properties defined by:
    category
    groupSpread
    complexityScore
    
seq     complexityScore category        genomicSpread   genomicRef      genomicRefType  genomicStart    genomicEnd      
genomicMismatch groupSpread     groupRef        groupRefAlign -groupRefType    groupStart      groupEnd        groupMismatch   groupMatchLength
        1       2       3       4       5       6       7       8

HCT ID, Lane, Sample ID, Cell Line, ER +/- Status, Reads
HCT20061, Lane 1, HS378, MCF7, +,  14,555,390 
HCT20062, Lane 2, HS379, MDA-MB-231, -,  15,147,037 
HCT20063, Lane 3, HS381, T47D, +,  14,701,254 
HCT20064, Lane 4, HS377, BT-20, -,  16,529,823 
HCT20065, Lane 5, HS375, BT-474, +,  16,443,361 
HCT20066, Lane 6, HS380, MDA-MB-468, -,  16,746,981 
HCT20067, Lane 7, HS382, ZR-75-1, +,  17,631,005 
HCT20068, Lane 8, HS376, MCF10A, normal,  19,155,607 
'''
import sys
import numpy as np

from collections import deque, defaultdict
from optparse import OptionParser

## Functions
class mirnaRecord:
    def __init__(self,atribute_list):
        
        self.seq=atribute_list[0]
        self.complexityScore=atribute_list[1]
        self.category=atribute_list[2]
        self.genomicSpread=atribute_list[3]
        self.genomicRef=atribute_list[4]
        self.genomicRefType=atribute_list[5]
        self.genomicStart=atribute_list[6]
        self.genomicEnd=atribute_list[7]
        self.genomicMismatch=atribute_list[8]
        self.groupSpread=atribute_list[9]
        self.groupRef=atribute_list[10]
        self.groupRefAlign=atribute_list[11]
        self.groupRefType=atribute_list[12]
        self.groupStart=atribute_list[13]
        self.groupEnd=atribute_list[14]
        self.groupMismatch=atribute_list[15]
        self.groupMatchLength=atribute_list[16]
        self.samples = np.array(atribute_list[17:])
        
    def get_all_atributes(self):
        myatributes = [self.seq, self.complexityScore, self.genomicSpread, self.genomicRef, self.genomicRefType, self.genomicStart, \
                       self.genomicEnd, self.genomicMismatch, self.groupSpread, self.groupRef, self.groupRefAlign, self.groupRefType,\
                       self.groupStart, self.groupEnd, self.groupMismatch, self.groupMatchLength]
        return myatributes
    
def parse_file(inputfile_name):
    """
    """
    inputfile = open(inputfile_name)
    record_list=deque()
    headers=[]
    for line in inputfile:
        fields = line.strip('\n\t\r').split('\t')
        
        if not headers and fields[0] == ('seq') and fields[1] == ('complexityScore'):
            header = fields
            headers=True
            continue
        if headers and len(fields)>1:
            record_list.append(mirnaRecord(fields))
        else:
            headers=False
    
    inputfile.close()
    
    return record_list, header

def get_length_mirs_dict(mirDB):
    """
    Get the length of the mirna sequence. 
    It can be from the mature mirna or the
    precursor mirna depending on mirDB
    """
    fastafile = open(mirDB)
    flag=False
    mirna_dict = defaultdict()
    for line in fastafile:
        #print line
        if line.startswith('>hsa') and not flag:
            fields = line.strip('>\n').split(' ') 
            flag=True
            continue
        
        if flag:
            mirna_dict[fields[0]] = len(line)            
            flag=False
    
    fastafile.close()
    
    return mirna_dict

    


def category(ip):
    catg = {'n':'none', 's':'short', 'b':'both', 'd':'gdna', 'g':'group'}[ip]
    return catg


## Main
if __name__ == '__main__':
    
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-t", "--category", dest="category",
                        help="category: ")
    optionparser.add_option("-p", "--genomicSpread", dest="genomicSpread",
                        help="int > 1")    
    
    '''
    optionparser.add_option("-d", "--complexityScore", dest="complexityScore",
                        help="float. No used yet")
    '''

    (options, args) = optionparser.parse_args()
    
    thcatg = category(options.category)
    inputfile_name = '/data/projects/iDEA/Small_RNA/compare_all8.txt'
    outputfile_name = '/data/projects/iDEA/Small_RNA/filtered_compare_all8.txt'
    qspec_mat_name = '/data/projects/iDEA/Small_RNA/qspecmat_compare_all8.txt'
    mirDB_filename = '/data/databases/miRBase/mature_hsa.fa'
    mirDBpre_filename = '/data/databases/miRBase/hairpin_hsa.fa'
    
    ###
    outputfile=open(outputfile_name,'w')
    qspec_mat_file = open('/data/projects/iDEA/Small_RNA/qspecmat_compare_all8.txt','w')
    qspec_header = ["protid","protLen","0","0","0","0","1","1","1"]
    qspec_mat_file.write(",".join(qspec_header).replace(',','\t')+'\n')
    # The average length of the miRNA should be obtained from the miRBase database. Use the mature or the precursor length.
    avg_length = 24
    avg_length_loop = 40
    sample_order=[0,2,4,6,1,3,5]
    
    # Get the mirna length. mature in this case
    mirna_dict = get_length_mirs_dict(mirDB_filename)
    mirna_pre_dict = get_length_mirs_dict(mirDBpre_filename)
    
    mirna_list, header = parse_file(inputfile_name)
    for i,thisrec in enumerate(mirna_list):
        if thisrec.category == thcatg and int(thisrec.genomicSpread) <= int(options.genomicSpread):
            try:
                mirna_length = mirna_dict[thisrec.groupRef]
            except KeyError:
                
                if thisrec.groupRefType == "pre":
                    #tmp_name = thisrec.groupRef.replace('mir','miR')
                    try:
                        mirna_length = mirna_pre_dict[thisrec.groupRef]
                        #mirna_length = mirna_dict[tmp_name]*3
                    except KeyError:
                        mirna_length = avg_length
                # Get the info from the loop miRBase file
                elif thisrec.groupRefType == "loop":
                    mirna_length = avg_length_loop
            
            outputfile.write(",".join(thisrec.get_all_atributes()).replace(',','\t')+'\n')
            qspec_line = [thisrec.seq+'_'+thisrec.groupRef,str(mirna_length)]
            qspec_mat_file.write(",".join(qspec_line).replace(",","\t")+'\t'+",".join(map(str,thisrec.samples[sample_order])).replace(",","\t")+'\n')
            
    outputfile.close()